(ns jna.win.psapi
  (:use jna
        jna.win.kernel32
        jna.win.nt
        jna.win.error))

(def-jna-struct PROCESS_MEMORY_COUNTERS
  [:DWORD  :cb]
  [:DWORD  :PageFaultCount]
  [:SIZE_T :PeakWorkingSetSize]
  [:SIZE_T :WorkingSetSize]
  [:SIZE_T :QuotaPeakPagedPoolUsage]
  [:SIZE_T :QuotaPagedPoolUsage]
  [:SIZE_T :QuotaPeakNonPagedPoolUsage]
  [:SIZE_T :QuotaNonPagedPoolUsage]
  [:SIZE_T :PagefileUsage]
  [:SIZE_T :PeakPagefileUsage])

(def-jna-struct PROCESS_MEMORY_COUNTERS_EX
  [:DWORD  :cb]
  [:DWORD  :PageFaultCount]
  [:SIZE_T :PeakWorkingSetSize]
  [:SIZE_T :WorkingSetSize]
  [:SIZE_T :QuotaPeakPagedPoolUsage]
  [:SIZE_T :QuotaPagedPoolUsage]
  [:SIZE_T :QuotaPeakNonPagedPoolUsage]
  [:SIZE_T :QuotaNonPagedPoolUsage]
  [:SIZE_T :PagefileUsage]
  [:SIZE_T :PeakPagefileUsage]
  [:SIZE_T :PrivateUsage])

(defjna Boolean psapi GetProcessMemoryInfo)
(defjna Boolean psapi QueryWorkingSet)

(defn- open-process [pid]
  (let [h (OpenProcess (bit-or PROCESS_QUERY_INFORMATION PROCESS_VM_READ) false (int pid))]
    (when (= com.sun.jna.Pointer/NULL h) (throw (Exception. "OpenProcess Fails"))) ; XXX: TODO: FIXME: check null
    h))

(defn- query-working-set
  ([h]
    (let [buffer (jna-array-new :SIZE_T (* 1024 128))
          b (QueryWorkingSet h (byref buffer) (size-of buffer))
          n (jna-array-get 0 buffer)]
      (if-not b
        (let [code (GetLastError)]
          (if (= code ERROR_BAD_LENGTH)
            (query-working-set h (jna-array-new :SIZE_T (inc 1)))
            (throw (Exception. (format "QueryWorkingSet fails. GetLastError: %s" code)))))
        buffer)))
  ([h buffer]
    (let [b (QueryWorkingSet h (byref buffer) (size-of buffer))]
      (if-not b
        (let [code (GetLastError)]
          (if (= code ERROR_BAD_LENGTH)
            (throw (Exception. "QueryWorkingSet fails: ERROR_BAD_LENGTH"))
            (throw (Exception. (format "QueryWorkingSet fails. GetLastError: %s" code)))))
        buffer))))

;(set! *warn-on-reflection* true)
(defmacro buffer-at [buffer i]
  (condp = os-bits
    64 `(.getLong ~buffer (* 8 ~i))
    32 `(.getInt  ~buffer (* 4 ~i))))

(defn- count-working-set-pages [n dWorkingSetPages]
  "Use Java version for better performance: mekit.internal.FastPieces/getPrivateWorkingSet"
  ; TODO: FIXME: very bad performance, how to fix it?
  ; FIXME: a simplified (and faster) version of http://www.codeproject.com/KB/cpp/XPWSPrivate.aspx, but incorrect in some cases
  (loop [i (int 1)
         dSharedPages 0]
    (if (<= i n)
      (let [page (buffer-at dWorkingSetPages i)
            dPageAddress (bit-and page 0xFFFFF000)]
        (if (or (< dPageAddress 0xC0000000) (< 0xE0000000 dPageAddress))
          (if-not (== 0 (bit-and page 0x100))
            (recur (unchecked-inc i) (unchecked-inc  dSharedPages))
            (recur (unchecked-inc i) dSharedPages))
          (recur (unchecked-inc i) dSharedPages)))
      (let [total (* n 4 1024)
            shared (* dSharedPages 4 1024)
            privates (- total shared)]
        [total privates shared])))) ; in KB

(defn get-process-working-set [pid] ; in bytes
  ; http://www.codeproject.com/KB/cpp/XPWSPrivate.aspx
  (let [h (open-process pid)]
    (try
      (let [dWorkingSetPages (query-working-set h)
            n (jna-array-get 0 dWorkingSetPages)
            pages (:buffer dWorkingSetPages)
         ;   [total privates shared] (count-working-set-pages n pages) ; bad performance, use Java
            [total privates shared] (mekit.internal.FastPieces/getPrivateWorkingSet n pages)]
        {:total   total
         :private privates
         :shared  shared})
      (finally
        (CloseHandle h)))))


(defn get-process-memory [pid]
  (let [h (open-process pid)
        counters (jna-struct-new PROCESS_MEMORY_COUNTERS_EX)]
    (GetProcessMemoryInfo h (:pointer counters) (jna-size-of counters))
    (try
      {:vm          (jna-struct-get :PrivateUsage   counters)
       :memory      (jna-struct-get :WorkingSetSize counters)
       :private     (jna-struct-get :PrivateUsage   counters)
       :working-set (jna-struct-get :WorkingSetSize counters)
       }
      (finally
        (CloseHandle h)))))

;(println (time (get-process-working-set 4132)))

